from ctypes import *
import httplib
import math
import platform
from time import time
import array

FHMQV_ECP_PRIVATE_KEY_LENGTH = 48
FHMQV_ECP_PUBLIC_KEY_LENGTH = 97
FHMQV_ECP_SHARED_KEY_LENGTH = 32

###############################################################################
# Util methods

def hexstr(s):
    return ' '.join(c.encode('hex') for c in s)

#FIXME: Investigate if the explicit copy is necessary here.
def ptr_to_bytestring(ptr, length):
    contents = []
    for i in range(0, length):
        contents.append(ptr[i])
    return array.array('B', contents).tostring()


###############################################################################
# Util classes

class GKSL_CLIENTHELLO(Structure):
    _fields_ = [("eph_public_key", c_ubyte * FHMQV_ECP_PUBLIC_KEY_LENGTH),
                ("eph_private_key", c_ubyte * FHMQV_ECP_PRIVATE_KEY_LENGTH),
                ("msg", c_ubyte * (FHMQV_ECP_PUBLIC_KEY_LENGTH + 32)),
                ("msg_len", c_size_t)]


class GKSL_SESSIONDATA(Structure):
    _fields_ = [("sid", c_ubyte * 8),
                ("thread_cnt", c_ubyte),
                ("next_free_iv_idx", c_ubyte),
                ("s2c_aes_key", c_ubyte * 32),
                ("c2s_aes_key", c_ubyte * 32),
                ("session_ivs", POINTER(c_ubyte))]


class GKSL_REQUEST(Structure):
    _fields_ = [("msg", POINTER(c_ubyte)),
                ("msg_len", c_size_t)]


class GKSL_RESPONSE(Structure):
    _fields_ = [("data", POINTER(c_ubyte)),
                ("data_len", c_size_t)]


class GKSL_TERMREQUEST(Structure):
    _fields_ = [("msg", c_ubyte * 48),
                ("msg_len", c_size_t)]


class GKSLError(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

###############################################################################
# Decorators

def connection_required(fn):
    def wrapper(*args):
        if args[0].conn is None:
            raise GKSLError('no connection')
        return fn(*args)

    return wrapper

def session_required(fn):
    def wrapper(*args):
        if args[0].session_data is None:
            raise GKSLError('no session')
        return fn(*args)

    return wrapper

def no_session_exists(fn):
    def wrapper(*args):
        if args[0].session_data is not None:
            raise GKSLError('session is already exists')
        return fn(*args)

    return wrapper

def no_connection_exists(fn):
    def wrapper(*args):
        if args[0].conn is not None:
            raise GKSLError('connection already exists')
        return fn(*args)

    return wrapper

###############################################################################

# load shared library libgksl
if platform.system() == 'Darwin':
    libgksl = CDLL("libgksl.dylib")
elif platform.system() == 'Windows':
    libgksl = CDLL("libgksl.dll")
else:
    libgksl = CDLL("libgksl.so")

libgksl.gksl_clienthello_init.restype = POINTER(GKSL_CLIENTHELLO)
libgksl.gksl_sessiondata_init.restype = POINTER(GKSL_SESSIONDATA)
libgksl.gksl_encode_request.restype = POINTER(GKSL_REQUEST)
libgksl.gksl_decode_response.restype = POINTER(GKSL_RESPONSE)
libgksl.gksl_termination_request.restype = POINTER(GKSL_TERMREQUEST)


class GKSLConnection(object):

    def __init__(self, host, port, thread_cnt, kpid, static_priv_key, static_pub_key, server_pub_key):
        self.host = host
        self.port = port
        self.thread_cnt = thread_cnt
        self.kpid = kpid
        self.static_priv_key = static_priv_key
        self.static_pub_key = static_pub_key
        self.server_pub_key = server_pub_key

        self.conn = None
        self.session_data = None

    def do_req(self, msg, path="/", verbose=False):
        self.conn.request("POST", path, msg)
        ts0 = time()
        resp = self.conn.getresponse()
        resp_msg = resp.read()
        ts1 = time()
        assert resp.status == 200

        if verbose:
            print ""
            print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<"
            print "{:.3f} ms | {} bytes out | {} bytes in".format((ts1 - ts0) * 1000, len(msg), len(resp_msg))
            print "Server response: ", hexstr(resp_msg)
            print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<"
            print ""

        return resp_msg

#############################################################################################################
### public API

    @no_connection_exists
    def connect(self):
        try:
            self.conn = httplib.HTTPConnection(self.host, self.port, True)
            self.conn.connect()
        except Exception as e:
            self.conn = None
            raise GKSLError('connection error')

    @connection_required
    @no_session_exists
    def init_session(self):
        def cleanup():
            libgksl.gksl_clienthello_free(gksl_ch_ptr)

        try:
            gksl_ch_ptr, gksl_sessiondata_ptr = None, None

            gksl_ch_ptr = libgksl.gksl_clienthello_init(self.thread_cnt, math.trunc(time()), 0)
            gksl_ch = gksl_ch_ptr.contents

            # send client hello message to the server
            resp = self.do_req(gksl_ch.msg)

            # initialize session data from the server response
            gksl_sessiondata_ptr = libgksl.gksl_sessiondata_init(resp, len(resp), self.thread_cnt,
                                                                 gksl_ch.eph_private_key, gksl_ch.eph_public_key,
                                                                 self.static_priv_key, self.static_pub_key,
                                                                 self.server_pub_key)

            if gksl_sessiondata_ptr is None:
                raise GKSLError('gksl session initialization error')

            self.session_data = gksl_sessiondata_ptr
            cleanup()
        except Exception as e:
            self.session_data = None
            cleanup()
            raise e

    @connection_required
    @session_required
    def do_app_req(self, msg):
        def cleanup():
            libgksl.gksl_request_free(gksl_app_req_ptr)
            libgksl.gksl_response_free(gksl_app_res_ptr)

        try:
            gksl_app_req_ptr, gksl_app_res_ptr = None, None

            gksl_app_req_ptr = libgksl.gksl_encode_request(self.session_data, msg, len(msg))
            gksl_app_req = ptr_to_bytestring(gksl_app_req_ptr.contents.msg, gksl_app_req_ptr.contents.msg_len)

            resp = self.do_req(gksl_app_req)

            gksl_app_res_ptr = libgksl.gksl_decode_response(self.session_data, resp, len(resp))
            res = ptr_to_bytestring(gksl_app_res_ptr.contents.data, gksl_app_res_ptr.contents.data_len)

            cleanup()
            return res
        except Exception as e:
            cleanup()
            raise e

    @connection_required
    @session_required
    def terminate_session(self):
        def cleanup():
            libgksl.gksl_request_free(gksl_term_req_ptr)
            libgksl.gksl_sessiondata_free(self.session_data)
            self.session_data = None

        try:
            gksl_term_req_ptr = libgksl.gksl_termination_request(self.session_data)

            resp = self.do_req(gksl_term_req_ptr.contents.msg)

            term_code = libgksl.gksl_check_termination_response(self.session_data, resp, len(resp))

            if term_code != 0:
                raise GKSLError('gksl session termination error')

            cleanup()
        except Exception as e:
            cleanup()
            raise e

    @connection_required
    def disconnect(self):
        self.conn.close()
        self.conn = None

###############################################################################
# EOF
